﻿#pragma once
#include<iostream>

using namespace std;
enum Color 
{
	RED,
	BLACK
};

template<class T, class K>
struct RedBlackTreeNode
{
	pair<T, K> _kv;
	RedBlackTreeNode<T, K>* _parent;
	RedBlackTreeNode<T, K>* _left;
	RedBlackTreeNode<T, K>* _right;
	Color _col;

	RedBlackTreeNode(const pair<T, K>& kv)
		:_kv(kv)
		,_parent(nullptr)
		,_left(nullptr)
		,_right(nullptr)
	{}
}; 

//红黑树的规则：
//1. 节点非黑及红
//2. 根节点是黑色的
//3. 如果⼀个结点是红色的，则它的两个孩⼦结点必须是黑色的，也就是说任意⼀条路径不会有连续的
//   红⾊结点。
//4. 对于任意⼀个结点，从该结点到其所有NULL结点的简单路径上，均包含相同数量的黑色结点

template<class T, class K>
class RedBlackTree
{
	typedef RedBlackTreeNode Node;
public:
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			Node* cur = new Node(kv);
			_root = cur;

			return true;
		}

		Node* cur = _root;
		Node* parent = nullptr;
		while (cur)
		{
			if (cur->_kv.first > kv.first)
			{
				parent = cur->_parent;
				cur = cur->_left;
			}
			else if (cur->_kv.first < kv.first)
			{
				parent = cur->_parent;
				cur = cur->_right;
			}
			else
			{
				return false;
			}
		}

		cur = new Node(kv);
		cur->_col = RED;
		if (parent->_kv.first > kv.first)
		{
			parent->_left = cur;
		}
		else if (parent->_kv.first < kv.first)
		{
			parent->_right = cur;
		}

		if (parent->_col == BLACK)
		{
			return true;
		}

		cur->_parent = parent;

		while (parent && parent->_col == RED)
		{
			Node* grandparent = parent->_parent;
			if (grandparent->_left == parent)
			{
				Node* uncle = grandparent->_right;
				//情况1：变色
				if (uncle && uncle->_col == RED)
				{
					parent->_col = unclde->_col = BLACK;
					grandparent->_col = RED;

					cur = grandparent;
					parent = cur->_parent;
				}
				else //情况2：变色+旋转
				{
					if (cur == parent->_left)
					{
						RotateR(grandparent);
						parent->_col = BLACK;
						cur->_col = parent->_col = RED;
					}
					else
					{
						RotateL(parent);
						RotateR(grandparent);
						cur->_col = BLACK;
						grandparent->_col = parent->_col = RED;
					}

					break;
				}
			}
			else if (grandparent->_right == parent)
			{
				Node* uncle = grandparent->_left;
				if (uncle && uncle->_col == RED)
				{
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 叔叔不存在，或者存在且为⿊
				{
					// 情况⼆：叔叔不存在或者存在且为⿊
					// 旋转+变⾊
					// g
					// u p
					// c
					if (cur == parent->_right)
					{
						RotateL(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else
					{
						// g
						// u p
						// c
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}


			_root->_col = BLACK;

			return true;

		}

	}

	Node* Find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first > key)
			{
				cur = cur->_left;
			}
			else if (cur->_kv.first < key)
			{
				cur = cur->_right;
			}
			else
			{
				return cur;
			}
		}

		return nullptr;
	}

private:
	Node* _root;
};